#pragma once
#include <array>
#include <cstdint>
#include <cstring>
#include <sys/mman.h>

/**
 *  redirect function call from api to stub.
 */
class AddrModifier {
  constexpr static size_t PAGE_ALIGN_BITS = 12;
  constexpr static size_t PAGE_SIZE = 1 << PAGE_ALIGN_BITS;
  constexpr static size_t JMP_INSTRUCTION_SIZE = 14;

public:
  AddrModifier(void* api, void* stub) : api_(api) {
    // update memory page attribute
    if (::mprotect(AlignToPageBoundary(api),
                   PAGE_SIZE * 2, PROT_EXEC | PROT_WRITE | PROT_READ) != 0) {
      return;
    }

    // backup origin asm code
    ::memcpy(backup_.data(), api, JMP_INSTRUCTION_SIZE);

    // get new asm code
    const auto buffer = GetJumpCode(stub);
    ::memcpy(api_, buffer.data(), JMP_INSTRUCTION_SIZE);
  }

  ~AddrModifier() {
    ::memcpy(api_, backup_.data(), JMP_INSTRUCTION_SIZE);
  }

  AddrModifier(const AddrModifier&) = delete;
  AddrModifier(const AddrModifier&&) = delete;
  AddrModifier& operator=(const AddrModifier&) = delete;
  AddrModifier& operator=(const AddrModifier&&) = delete;

  void* GetMethod() const {
    return method_;
  }

  void SetMethod(void* method) {
    method_ = method;
  }

  static std::array<uint8_t, AddrModifier::JMP_INSTRUCTION_SIZE> GetJumpCode(
      void* dest) {
    // asm: jmp qword ptr [rip]
    std::array<uint8_t, JMP_INSTRUCTION_SIZE> buffer{
        0xFF, // 多功能操作码
        0x25, //  ModR/M 字节，表示 jmp r/m，即间接跳转
        0x00, 0x00, 0x00, 0x00, // 32位偏移地址
        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // ptr
    };

    *reinterpret_cast<uintptr_t*>(buffer.data() + 6) = reinterpret_cast<
      uintptr_t>(dest);
    return buffer;
  }

  static void* AlignToPageBoundary(const void* addr) {
    return reinterpret_cast<void*>(
      reinterpret_cast<uintptr_t>(addr) & (~((1 << (PAGE_ALIGN_BITS)) - 1)));
  }

private:
  void* const api_ = nullptr;
  void* method_ = nullptr;
  std::array<uint8_t, JMP_INSTRUCTION_SIZE> backup_{};
};